#include <stdio.h>
#include <stdlib.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/ioctl.h>
#include <string.h>
#include <pthread.h>
#include <time.h>
#include <errno.h>
#include <signal.h>

#define DEVICE_PATH "/dev/pcie-rkep-0000:09:00.0"
#define PCIE_EP_DMA_TEST _IOW('P', 7, struct pcie_ep_dma_test)

enum dma_dir {
    DMA_FROM_BUS,
    DMA_TO_BUS,
};

struct pcie_ep_dma_test {
    u_int32_t direction;
    u_int32_t chn;
    u_int32_t size;
    u_int32_t timeout_ms;
};

// 线程同步结构
struct dma_controller {
    int fd;
    int direction;
    int channel;
    volatile int running;
    pthread_mutex_t lock;

    // 统计信息
    unsigned long total_bytes;
    unsigned long total_time_ns;
    unsigned test_count;
    unsigned error_count;
};

// 获取当前时间（纳秒）
static inline u_int64_t get_time_ns() {
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return (u_int64_t)ts.tv_sec * 1000000000ULL + (u_int64_t)ts.tv_nsec;
}

// DMA测试线程函数
void* dma_test_thread(void *arg) {
    struct dma_controller *ctrl = (struct dma_controller *)arg;
    struct pcie_ep_dma_test test_args = {
        .direction = ctrl->direction,
        .chn = ctrl->channel,
        .size = 2*1024*1024, // 2MB
        .timeout_ms = 100       // 无限等待
    };

    while (ctrl->running) {
        pthread_mutex_lock(&ctrl->lock);

        u_int64_t start_time = get_time_ns();
        int ret = ioctl(ctrl->fd, PCIE_EP_DMA_TEST, &test_args);
        u_int64_t end_time = get_time_ns();

        if (ret == 0) {
            ctrl->total_bytes += test_args.size;
            ctrl->total_time_ns += (end_time - start_time);
            ctrl->test_count++;
        } else {
            ctrl->error_count++;
            fprintf(stderr, "DMA test failed: %s\n", strerror(errno));
        }

        pthread_mutex_unlock(&ctrl->lock);

        // 短暂休眠避免CPU占用过高
        usleep(1000); // 1ms
    }

    return NULL;
}

// 信号处理函数
volatile sig_atomic_t stop_request = 0;
void sigint_handler(int sig) {
    stop_request = 1;
}

int main(int argc, char *argv[]) {
    if (argc != 4) {
        printf("Usage: %s <direction> <channel> <test_duration_sec>\n", argv[0]);
        printf("  direction: 1=TO_BUS, 0=FROM_BUS\n");
        printf("  channel:   0 or 1\n");
        printf("  test_duration_sec: Test duration in seconds\n");
        return 1;
    }

    int direction = atoi(argv[1]);
    int channel = atoi(argv[2]);
    int test_duration = atoi(argv[3]);

    if (direction != 0 && direction != 1) {
        printf("Invalid direction: must be 0 or 1\n");
        return 1;
    }

    if (channel != 0 && channel != 1) {
        printf("Invalid channel: must be 0 or 1\n");
        return 1;
    }

    if (test_duration <= 0) {
        printf("Test duration must be positive integer\n");
        return 1;
    }

    int fd = open(DEVICE_PATH, O_RDWR);
    if (fd < 0) {
        perror("Failed to open device");
        return 1;
    }

    // 初始化控制器
    struct dma_controller ctrl = {
        .fd = fd,
        .direction = direction,
        .channel = channel,
        .running = 1,
        .total_bytes = 0,
        .total_time_ns = 0,
        .test_count = 0,
        .error_count = 0
    };
    pthread_mutex_init(&ctrl.lock, NULL);

    // 注册信号处理
    signal(SIGINT, sigint_handler);
    
    // 创建测试线程
    pthread_t test_thread;
    if (pthread_create(&test_thread, NULL, dma_test_thread, &ctrl)) {
        perror("Failed to create test thread");
        close(fd);
        return 1;
    }

    printf("Starting DMA test for %d seconds... (Press Ctrl+C to stop early)\n", test_duration);

    // 实时带宽统计
    time_t start_time = time(NULL);
    time_t last_print = start_time;
    unsigned long last_bytes = 0;

    while (!stop_request) {
        sleep(1);
        time_t current_time = time(NULL);

        // 每秒钟打印实时带宽
        if (current_time > last_print) {
            pthread_mutex_lock(&ctrl.lock);
            unsigned long current_bytes = ctrl.total_bytes;
            pthread_mutex_unlock(&ctrl.lock);

            double elapsed = difftime(current_time, last_print);
            double bandwidth = (elapsed > 0) ?
                (double)(current_bytes - last_bytes) / (1024.0 * 1024.0) / elapsed : 0;

            printf("Real-time bandwidth: %.2f MB/s (Total: %.2f MB)\n",
                   bandwidth, (double)current_bytes / (1024 * 1024));

            last_bytes = current_bytes;
            last_print = current_time;
        }

        // 检查是否达到测试时间
        if (difftime(current_time, start_time) >= test_duration) {
            break;
        }
    }

    // 停止测试线程
    ctrl.running = 0;
    pthread_join(test_thread, NULL);

    // 计算结果
    double total_time_sec = (double)ctrl.total_time_ns / 1e9;
    double bandwidth = (total_time_sec > 0) ?
        (double)ctrl.total_bytes / (1024.0 * 1024.0) / total_time_sec : 0;

    // 打印最终结果
    printf("\n===== Test Summary =====\n");
    printf("  Test duration:    %d seconds\n", test_duration);
    printf("  Total transfers:  %u\n", ctrl.test_count);
    printf("  Failed transfers: %u\n", ctrl.error_count);
    printf("  Total data:       %.2f MB\n", (double)ctrl.total_bytes / (1024 * 1024));
    printf("  Total time:       %.6f seconds\n", total_time_sec);
    printf("  Average bandwidth: %.2f MB/s\n", bandwidth);

    // 清理资源
    pthread_mutex_destroy(&ctrl.lock);
    close(fd);
    return 0;
}
